#!/usr/bin/pytohn3
# -*- coding:UTF-8 -*-

# __author__ : xgzhang
# __create_time__   : 2021-08-21-18:10
# __useful_modify__ : 2021-08-23-09:00-11：45
# __useful_modify__ : 2021-08-24-15:00-18：30
# __useful_modify__ : 2021-08-25-14:00-18：45
# __other_time__    : computer boom / always booming

# 导入scapy包
import scapy
import scapy.all
from scapy.all import *
from scapy.layers.dhcp import BOOTP, DHCP
from scapy.layers.inet import IP, UDP
from scapy.layers.l2 import Ether
import binascii

# 多线程相关包
import threading
from concurrent.futures import ThreadPoolExecutor

# 队列相关包
import queue
from queue import Queue
import transitions

# 时间包
import datetime
import copy


# 定义常量：网卡
INFACE_NAME = 'en1'

# 定义所有状态的列表
states_lst = [
    'no',
    'discover',
    'offer',
    'request',
    'ack',
    'release',
    'close',
    'req_renewal',
]

# 定义状态切换器
# 也就是当发生什么时从哪个状态转换到哪个状态
transitions_lst = [
    ['init', 'no', 'discover'],
    ['dis_off', 'discover', 'offer'],
    ['off_req', 'offer', 'request'],
    ['req_ack', 'request', 'ack'],
    ['to_close', '*', 'close'],
    ['renewal_discover', 'ack', 'discover'],
    ['renewal_request', 'ack', 'request']
]

# 创建互斥锁,默认不上锁
lock = threading.Lock()


# 变换状态，进行数据包处理中转
def change_state(xid):
    global xid_client_dict
    global send_queue

    client = xid_client_dict[xid]
    print(client.state)
    if client.state == 'no':
        client.init()
        send_queue.put([client.macList, 'discover'])

    elif client.state == 'discover':
        client.dis_off()

    elif client.state == 'offer':
        print("offer包转换状态")
        client.off_req()
        send_queue.put([client.macList, 'request'])
        print("request 包放入队列")

    elif client.state == 'request':
        client.req_ack()

    elif client.state == 'ack':
        # xid_client_dict[xid].renewalTime =
        print("租约时长 ： " + str(xid_client_dict[xid].packet[DHCP].options[2][1]))
        # 存放租约时长
        lock.acquire()
        xid_client_dict[xid].renewalTimeList[0] = xid_client_dict[xid].packet[DHCP].options[2][1]
        lock.release()
        # 存放获取到IP和租约的起始时间
        lock.acquire()
        xid_client_dict[xid].renewalTimeList[1] = datetime.datetime.now()
        lock.release()
        print("四个阶段完成，申请到IP")

    # 申请续约不需要在此处理


## **********************************  初始化client，根据初始请求状态，改变状态，将初始化请求数据压入发送队列 ******************************************************

# 随机生成列表
# 初始化状态state为：no
# 返回列表 ： [xid, macaddress,client_mac_id,req,states]
def get_rand_mac():
    randomReq = ['add', 'close']
    macList = []
    # 列表：[xid, macaddress,client_mac_id,req,states]
    xid_random = random.randint(1, 900000000)
    mac_random = str(RandMAC())
    client_mac_id = binascii.unhexlify(mac_random.replace(':', ''))
    req_name = randomReq[random.randint(0, 1)]

    macList.append(xid_random)
    macList.append(mac_random)
    macList.append(client_mac_id)
    macList.append(req_name)
    macList.append('no')

    return macList


# 初始化client，生成状态控制机器
# client = Client(macList, packet)
# 所以client有两个属性值：随机生成的列表，初始化空的数据包
# 返回client
def init_client(macList):
    ethernet_discover = Ether()
    ip_discover = IP()
    udp_discover = UDP()
    bootp_discover = BOOTP()
    dhcp_discover = DHCP()
    packet = ethernet_discover / ip_discover / udp_discover / bootp_discover / dhcp_discover

    client = Client(macList, packet)

    # 生成一个状态机控制器
    machine = transitions.Machine(model=client,
                                  states=states_lst,
                                  transitions=transitions_lst,
                                  initial='no'
                                  )

    return client


# 对随机生成的请求进行处理
# add / close
def handle_init():
    global xid_client_dict
    global xid_client_dict_bak_to_check_renewal

    macList = get_rand_mac()
    # print("macList: ")
    # print(macList)

    if macList[3] == 'add':
        client = init_client(macList)
        # print("client: ")
        # print(client)
        lock.acquire()
        xid_client_dict[macList[0]] = client
        xid_client_dict_bak_to_check_renewal[macList[0]] = client
        lock.release()
        # print("xid_client_dict: ")
        # print(xid_client_dict)
        change_state(macList[0])
    if macList[3] == 'close':
        if macList[0] in xid_client_dict:
            if xid_client_dict[macList[0]].state == 'ack':
                print("发送关闭请求")
                change_state(xid_client_dict[macList[0]])
            else:
                print("终端IP正在申请中。。。。，请稍后申请关闭")
        else:
            print("终端未申请IP，不需要关闭")


## **********************************  从发送队列取出数据，进行处理，判断需要发送的包，模拟数据包并发送 ***************************************

def build_discover_packet(getQueueEle):
    global temp_dis_state_dict

    print("hahahahahahahahahahahahhaahhaahhh")
    # 模拟discover数据包
    ethernet_discover = Ether(src=getQueueEle[1], dst='ff:ff:ff:ff:ff:ff')
    ip_discover = IP(src='0.0.0.0', dst='255.255.255.255')
    udp_discover = UDP(sport=68, dport=67)
    bootp_discover = BOOTP(chaddr=getQueueEle[2], xid=getQueueEle[0])
    dhcp_discover = DHCP(options=[("message-type", "discover"), "end"])
    packet_discover = ethernet_discover / ip_discover / udp_discover / bootp_discover / dhcp_discover

    # 发送discover数据包
    lock.acquire()
    sendp(packet_discover, iface=INFACE_NAME)
    print("已发送discover包 ： " + str(getQueueEle[0]))
    lock.release()

    if temp_dis_state_dict[getQueueEle[0]][1] == 0:
        change_state(getQueueEle[0])
        lock.acquire()
        temp_dis_state_dict[getQueueEle[0]] = ["discover", 1, datetime.datetime.now()]
        lock.release()
    else:
        lock.acquire()
        temp_dis_state_dict[getQueueEle[0]][1] = temp_dis_state_dict[getQueueEle[0]][1] + 1
        lock.release()


def build_request_packet(macList):
    global temp_req_state_dict
    global xid_client_dict
    global xid_client_dict_bak_to_check_renewal

    # 模拟request包
    packet = xid_client_dict[macList[0]].packet
    ethernet_request = Ether(src=packet[Ether].dst, dst="ff:ff:ff:ff:ff:ff")
    ip_request = IP(src='0.0.0.0', dst='255.255.255.255')
    udp_request = UDP(sport=68, dport=67)
    bootp_request = BOOTP(chaddr=packet[BOOTP].chaddr, xid=packet[BOOTP].xid)
    dhcp_request = DHCP(options=[("message-type", 'request'), ("server_id", packet[DHCP].options[1][1]),
                                 ("requested_addr", packet[BOOTP].yiaddr), "end"])
    packet_request = ethernet_request / ip_request / udp_request / bootp_request / dhcp_request

    # 发送request数据包
    lock.acquire()
    sendp(packet_request, iface=INFACE_NAME)
    print("已发送request包 ：" + str(macList[0]))
    lock.release()

    lock.acquire()
    xid_client_dict[packet[BOOTP].xid].packet = packet_request
    xid_client_dict_bak_to_check_renewal[packet[BOOTP].xid].packet = packet_request
    lock.release()

    if temp_req_state_dict[packet[BOOTP].xid][1] == 0:
        change_state(packet[BOOTP].xid)
        lock.acquire()
        temp_req_state_dict[packet[BOOTP].xid] = ["request", 1, datetime.datetime.now()]
        lock.release()
    else:
        lock.acquire()
        temp_req_state_dict[packet[BOOTP].xid][1] = temp_req_state_dict[packet[BOOTP].xid][1] + 1
        lock.release()


def build_request_renewal_packet(macList):
    global xid_client_dict
    global temp_renewal_state_dict

    # 模拟request包
    packet = xid_client_dict[macList[0]].packet
    ethernet_request = Ether(src=packet[Ether].dst, dst=packet[Ether].src)
    ip_request = IP(src=packet[IP].dst, dst=packet[IP].src)
    udp_request = UDP(sport=68, dport=67)
    bootp_request = BOOTP(chaddr=packet[BOOTP].chaddr, xid=packet[BOOTP].xid)
    dhcp_request = DHCP(options=[("message-type", 'request'), ("server_id", packet[DHCP].options[1][1]),
                                 ("requested_addr", packet[BOOTP].yiaddr), "end"])
    packet_request = ethernet_request / ip_request / udp_request / bootp_request / dhcp_request

    # 发送request数据包
    lock.acquire()
    sendp(packet_request, iface=INFACE_NAME)
    print("已发送renewal request包 ：" + str(macList[0]))
    lock.release()
    # xid_client_dict[macList[0]].packet = packet_request

    if temp_renewal_state_dict[packet[BOOTP].xid][1] == 0:
        lock.acquire()
        temp_renewal_state_dict[packet[BOOTP].xid] = ["request", 1, datetime.datetime.now()]
        lock.release()
    else:
        lock.acquire()
        temp_renewal_state_dict[packet[BOOTP].xid][1] = temp_renewal_state_dict[packet[BOOTP].xid][1] + 1
        lock.release()


def build_release_packet(macList):
    global xid_client_dict
    global xid_client_dict_bak_to_check_renewal
    global temp_release_state_dict

    # 模拟releaset包
    packet = xid_client_dict[macList[0]].packet
    ethernet_release = Ether(src=packet[Ether].dst, dst=packet[Ether].src)
    ip_release = IP(src=packet[IP].dst, dst=packet[IP].src)
    udp_release = UDP(sport=68, dport=67)
    bootp_release = BOOTP(chaddr=packet[BOOTP].chaddr, xid=packet[BOOTP].xid)
    dhcp_release = DHCP(options=[("message-type", 'release'), ("server_id", packet[DHCP].options[1][1]),
                                 ("requested_addr", packet[BOOTP].yiaddr), "end"])
    packet_release = ethernet_release / ip_release / udp_release / bootp_release / dhcp_release

    # 发送release数据包
    lock.acquire()
    sendp(packet_release, iface=INFACE_NAME)
    print("已发送request包 ：" + str(macList[0]))
    lock.release()

    if temp_release_state_dict[packet[BOOTP].xid][1] == 0:
        lock.acquire()
        temp_release_state_dict[packet[BOOTP].xid] = ["request", 1, datetime.datetime.now()]
        lock.release()
        lock.acquire()
        xid_client_dict[macList[0]].packet = packet_release
        xid_client_dict_bak_to_check_renewal[macList[0]].packet = packet_release
        lock.release()
        # print(packet_release)
    else:
        lock.acquire()
        temp_release_state_dict[packet[BOOTP].xid][1] = temp_release_state_dict[packet[BOOTP].xid][1] + 1
        lock.release()


# 初始化线程，在线程池提交的函数：主要发送数据包，上面四种
def send_packet():
    global send_queue
    global send_queue
    global xid_client_dict
    global xid_client_dict_bak_to_check_renewal
    global temp_dis_state_dict
    global temp_req_state_dict
    global temp_renewal_state_dict
    global temp_release_state_dict

    client_data = send_queue.get()
    # [[xid, macaddress,client_mac_id,req,states],"state"]

    #  发送discover包
    if client_data[1] == 'discover':
        if client_data[0][0] not in temp_dis_state_dict:
            lock.acquire()
            temp_dis_state_dict[client_data[0][0]] = ["discover", 0, datetime.datetime.now()]
            lock.release()
            print("判断发送discover包，第一次发送discover包")
            build_discover_packet(client_data[0])
        else:
            print("重复发送discover包")
            build_discover_packet(client_data[0])

    # 判断发送request包
    if client_data[1] == 'request':
        if client_data[0][0] not in temp_req_state_dict:
            lock.acquire()
            temp_req_state_dict[client_data[0][0]] = ["request", 0, datetime.datetime.now()]
            lock.release()
            print("判断发送request包")
            build_request_packet(client_data[0])
        else:
            print("重复发送request包")
            build_request_packet(client_data[0])

    # 判断发送释放关闭release包
    if client_data[1] == 'release':
        if client_data[0][0] not in temp_release_state_dict:
            lock.acquire()
            temp_release_state_dict[client_data[0][0]] = ["release", 0, datetime.datetime.now()]
            lock.release()
            print("判断发送release包")
            build_release_packet(client_data[0])
        else:
            print("重复发送release包")
            build_release_packet(client_data[0])

    # 判断发送续租的包
    if client_data[1] == 'renewal':
        if client_data[0][0] not in temp_renewal_state_dict:
            lock.acquire()
            temp_renewal_state_dict[client_data[0][0]] = ["release", 0, datetime.datetime.now()]
            lock.release()
            print("判断发送renewal包")
            build_request_renewal_packet(client_data[0])
        else:
            print("重复发送renewal包")
            build_request_renewal_packet(client_data[0])


## **********************************  sniff_packet()嗅探数据包，并将数据包压入接收队列  *******************************************************
# 嗅探完数据，进行处理，压入队列，嗅探主程序sniff回调的函数
def put_packet_queue(packet):
    global receive_queue
    global xid_client_dict
    if DHCP in packet and packet[BOOTP].xid in xid_client_dict:
        if packet[DHCP].options[0][1] == 2:
            receive_queue_element = [packet[BOOTP].xid, "offer", packet]
            receive_queue.put(receive_queue_element)
            print("queue offer 大小 ： " + str(receive_queue.qsize()) + "  id: " + str(packet[BOOTP].xid))
            # print(receive_queue.get())
        elif packet[DHCP].options[0][1] == 5:
            receive_queue_element = [packet[BOOTP].xid, "ack", packet]
            receive_queue.put(receive_queue_element)
            print("queue ack 大小 ： " + str(receive_queue.qsize()) + "  id: " + str(packet[BOOTP].xid))
        elif packet[DHCP].options[0][1] == 6:
            if packet[BOOTP].xid in renewal_list:
                build_discover_packet(xid_client_dict[packet[BOOTP].xid].macList)


# 初始化线程：嗅探主程序
def sniff_packet():
    sniff(filter="udp and (port 67 or 68)", iface=INFACE_NAME, prn=put_packet_queue)


## **********************************  从接收队列取出数据，并处理，更改xid_client_dict，调用change_state（）处理中转  ****************************

def handle_receive_queue():
    global receive_queue
    global xid_client_dict
    global xid_client_dict_bak_to_check_renewal
    global renewal_list
    global temp_dis_state_dict

    # 接收队列的每个元素是个列表： [xid, 状态，packet]
    receive_queue_get_ele = receive_queue.get()

    # 覆盖更新xid_client_dict中的packet包
    lock.acquire()
    xid_client_dict[receive_queue_get_ele[0]].packet = receive_queue_get_ele[2]
    xid_client_dict_bak_to_check_renewal[receive_queue_get_ele[0]].packet = receive_queue_get_ele[2]
    lock.release()

    if receive_queue_get_ele[1] == "offer":
        # 如果还在处理这个包的超时判断中，那就处理，否则说明此包已经处理完毕，那这个队列弹出来的就不需要处理，丢掉即可
        if receive_queue_get_ele[0] in temp_dis_state_dict:
            if (datetime.datetime.now() - temp_dis_state_dict[receive_queue_get_ele[0]][2]).total_seconds() < 3:
                print("未超时，dis --> req")
                change_state(receive_queue_get_ele[0])
            elif temp_dis_state_dict[receive_queue_get_ele[0]][1] < 3:
                print("超时，但未超过三次，继续发送discover包")
                build_discover_packet(xid_client_dict[receive_queue_get_ele[0]].macList)
            else:
                print("已重复发送discover包三次，申请IP失败")
                lock.acquire()
                del temp_dis_state_dict[receive_queue_get_ele[0]]
                lock.release()
                lock.acquire()
                del xid_client_dict[receive_queue_get_ele[0]]
                del xid_client_dict_bak_to_check_renewal[receive_queue_get_ele[0]]
                lock.release()

    if receive_queue_get_ele[1] == "ack":
        # 判断request回来的包
        if receive_queue_get_ele[0] in temp_req_state_dict:
            if (datetime.datetime.now() - temp_req_state_dict[receive_queue_get_ele[0]][2]).total_seconds() < 3:
                print("未超时，req--> ack")
                change_state(receive_queue_get_ele[0])
            elif temp_req_state_dict[receive_queue_get_ele[0]][1] < 3:
                print("超时，但未超过三次，继续发送request包")
                build_request_packet(xid_client_dict[receive_queue_get_ele[0]].macList)
            else:
                print("已重复发送request包三次，申请IP失败")
                lock.acquire()
                del temp_req_state_dict[receive_queue_get_ele[0]]
                lock.release()
                lock.acquire()
                del xid_client_dict[receive_queue_get_ele[0]]
                del xid_client_dict_bak_to_check_renewal[receive_queue_get_ele[0]]
                lock.release()

        # 判断release回来的包
        if receive_queue_get_ele[0] in temp_release_state_dict:
            if (datetime.datetime.now() - temp_release_state_dict[receive_queue_get_ele[0]][2]).total_seconds() < 3:
                print("未超时，release --> close")
                xid_client_dict[receive_queue_get_ele[0]].stop()
            elif temp_release_state_dict[receive_queue_get_ele[0]][1] < 3:
                print("超时，但未超过三次，继续发送release包")
                build_release_packet(xid_client_dict[receive_queue_get_ele[0]].macList)
            else:
                print("已重复发送release包三次，释放关闭IP失败")
                lock.acquire()
                del temp_release_state_dict[receive_queue_get_ele[0]]
                lock.release()
                lock.acquire()
                del xid_client_dict[receive_queue_get_ele[0]]
                del xid_client_dict_bak_to_check_renewal[receive_queue_get_ele[0]]
                lock.release()
                # 销毁状态机
                xid_client_dict[receive_queue_get_ele[0]].stop()

        # 判断renewal回来的包
        if receive_queue_get_ele[0] in temp_renewal_state_dict:
            if (datetime.datetime.now() - temp_renewal_state_dict[receive_queue_get_ele[0]][2]).total_seconds() < 3:
                print("未超时，ack --> renewal")
                print("续约成功")
                # 重置此xid的租约时长和起始时间
                lock.acquire()
                xid_client_dict[receive_queue_get_ele[0]].renewalTimeList[0] = \
                    receive_queue_get_ele[2][DHCP].options[2][1]
                lock.release()
                lock.acquire()
                xid_client_dict[receive_queue_get_ele[0]].renewalTimeList[1] = datetime.datetime.now()
                lock.release()

                # 续约列表删除此xid
                lock.acquire()
                renewal_list.remove(receive_queue_get_ele[0])
                lock.release()
            elif temp_renewal_state_dict[receive_queue_get_ele[0]][1] < 3:
                print("超时，但未超过三次，继续发送renewal包")
                build_request_renewal_packet(xid_client_dict[receive_queue_get_ele[0]].macList)
            else:
                print("已重复发送renewal包三次，续约IP失败")
                lock.acquire()
                del temp_renewal_state_dict[receive_queue_get_ele[0]]
                lock.release()
                lock.acquire()
                del xid_client_dict[receive_queue_get_ele[0]]
                del xid_client_dict_bak_to_check_renewal[receive_queue_get_ele[0]]
                lock.release()


## **********************************  检测xid_client_dict ,判断续约    ************************************************************ ***************************

def check_signal_ele(k):
    global send_queue
    global renewal_list
    global xid_client_dict
    global xid_client_dict_bak_to_check_renewal

    if (datetime.datetime.now() - xid_client_dict_bak_to_check_renewal[k].renewalTimeList[1]).seconds <= (
            xid_client_dict_bak_to_check_renewal[k].renewalTimeList[0] * 0.5):
        if (datetime.datetime.now() - xid_client_dict_bak_to_check_renewal[k].renewalTimeList[1]).seconds <= (
                xid_client_dict_bak_to_check_renewal[k].renewalTimeList[0] * 0.75):
            # 续约
            send_queue.put([xid_client_dict_bak_to_check_renewal[k].macList, 'discover'])
            lock.acquire()
            xid_client_dict[k].renewal_discover()
            lock.release()
        # 续约
        send_queue.put([xid_client_dict_bak_to_check_renewal[k].macList, 'renewal'])
        lock.acquire()
        xid_client_dict[k].renewal_request()
        lock.release()
        lock.acquire()
        renewal_list.append(k)
        lock.release()


## **********************************  classes  ****************************************************************************************************************************************************************


class Client():
    def __init__(self, macList, packet):
        # self.xid_random = macList[0]
        # self.mac_random = macList[1]
        # self.client_mac_id = macList[2]
        # self.req_name = macList[3]

        # 初始化生成的列表：[xid, macaddress,client_mac_id,req,states]
        self.macList = macList
        # 初始化生成的空的数据包
        self.packet = packet
        # 初始化生成时间字典：{sate:[now_time, next_time, time_minus_val]}
        self.timeList = {}
        # 初始化租约时长:默认为秒，一天86400秒
        self.renewalTimeList = [0, datetime.datetime.now()]


class InitClient(threading.Thread):
    def __init__(self, name, send_queue):
        threading.Thread.__init__(self, name=name)
        self.send_queue = send_queue

    def run(self):
        with ThreadPoolExecutor(max_workers=1) as handle_init_thread:
            for i in range(200):
                handle_init_thread.submit(handle_init)


class SendQueue_right(threading.Thread):
    def __init__(self, name, send_queue):
        threading.Thread.__init__(self, name=name)
        self.send_queue = send_queue

    def run(self):
        with ThreadPoolExecutor(max_workers=1) as handle_send_thread:
            while (True):
                handle_send_thread.submit(send_packet)


class Handle_receive_queue(threading.Thread):
    def __init__(self, name, receive_queue):
        threading.Thread.__init__(self, name=name)
        self.receive_queue = receive_queue

    def run(self):
        with ThreadPoolExecutor(max_workers=1) as handle_receive_thread:
            while (True):
                handle_receive_thread.submit(handle_receive_queue)


class Check_Renewal(threading.Thread):
    def __init__(self, name, xid_client_dict_bak_to_check_renewal):
        threading.Thread.__init__(self, name=name)
        self.xid_client_dict_bak_to_check_renewal = xid_client_dict_bak_to_check_renewal

    def run(self):
        with ThreadPoolExecutor(max_workers=1) as check_renewal_thread:
            while True:
                it = xid_client_dict_bak_to_check_renewal.__iter__()
                while True:
                    try:
                        k = it.__next__()
                        check_renewal_thread.submit(check_signal_ele, k)
                    except:
                        break

if __name__ == '__main__':
    # 保存初始化client，格式为{xid : client}  /// 最要的*所有的数据
    # global xid_client_dict
    xid_client_dict = {}

    # global xid_client_dict_bak_to_check_renewal
    xid_client_dict_bak_to_check_renewal = {}

    # 暂存 discover 当前状态的当前时间
    # global temp_dis_state_dict
    temp_dis_state_dict = {}

    # 暂存 request 当前状态的当前时间
    # global temp_req_state_dict
    temp_req_state_dict = {}

    # 暂存 release 当前状态的时间
    # global temp_release_state_dict
    temp_release_state_dict = {}

    # 暂存 renewal 当前状态时间
    # global temp_renewal_state_dict
    temp_renewal_state_dict = {}

    # 暂时存储需要续约的xid
    # global renewal_list
    renewal_list = []

    # 创建发送队列
    # global send_queue
    send_queue = Queue()

    # 创建接收队列
    # global receive_queue
    receive_queue = queue.Queue()

    # 嗅探进程，入接收队列
    sniff_thread = threading.Thread(target=sniff_packet)
    sniff_thread.start()

    # 初始化请求进程，入发送队列
    init_client_thread = InitClient("initclient", send_queue)
    init_client_thread.start()

    # 巡检租约超时进程
    # check_renewal_time_thread = Check_Renewal("check_renewal", xid_client_dict_bak_to_check_renewal)
    # check_renewal_time_thread.start()

    # 处理接收队列进程，拿出数据消费
    handle_receive_queue_thread = Handle_receive_queue("receive_consumer", receive_queue)
    handle_receive_queue_thread.start()

    # 处理发送队列，发送数据消费
    send_queue_thread = SendQueue_right("send_consumer", send_queue)
    send_queue_thread.start()
